///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/energy_acc.h
///
/// \brief   header file for the energy accumulator class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_FQMC_ENERGY_ACC_HEADER
#define FORMIC_FQMC_ENERGY_ACC_HEADER

#include<complex>
#include<vector>
#include<string>
#include<numeric>
#include<cassert>
#include<algorithm>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/fqmc/accumulator.h>
#include<formic/wfn/configuration.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/output/output.h>
#include<formic/exception.h>
#include<formic/lapack/interface.h>
#include<formic/random/random.h>
#include<formic/vector/vector.h>

namespace formic {

  namespace fqmc {

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  a class to accumulate energies during a markov chain evolution
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    template <class S> class EnAcc : public formic::Accumulator<S> {

      private:

        /// \brief  pointer to the guiding function
        boost::shared_ptr<formic::CoeffObj<S> > _gf;

        /// \brief  pointer to the value function
        boost::shared_ptr<formic::CoeffObj<S> > _vf;

        /// \brief  whether exact sampling is being used
        bool _exact_sampling;

        /// \brief  whether to print local energies as they are sampled
        bool _print_local_e_during_sampling;

        /// \brief  energy
        double _energy;

        /// \brief  variance
        double _variance;

        /// \brief  estimated statistical error
        double _serr;

        /// \brief  a history of the sampled probabilities
        std::vector<double> _history_p;

        /// \brief  a history of the |value/guiding|^2 ratios
        std::vector<double> _history_v;

        /// \brief  a history of the sampled local energies
        std::vector<double> _history_e;

        /// \brief  a history of the sampled local energies times the |value/guiding|^2 ratios
        std::vector<double> _history_ev;

        /// \brief  the maximum lenth for which to compute autocorrelations
        int _max_ac_length;

      public:

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  constructs the accumulator around a specific coefficient object
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        EnAcc(const boost::shared_ptr<formic::CoeffObj<S> > & gf,
              const boost::shared_ptr<formic::CoeffObj<S> > & vf,
              const int max_ac_length,
              const bool exact_sampling,
              const bool print_local_e_during_sampling)
          : _gf(gf),
            _vf(vf),
            _exact_sampling(exact_sampling),
            _print_local_e_during_sampling(print_local_e_during_sampling),
            _energy(0.0),
            _variance(0.0),
            _serr(-1.0),
            _history_p(),
            _history_v(),
            _history_e(),
            _history_ev(),
            _max_ac_length(max_ac_length)
        {}

        /// \brief  return the accumulator's name
        std::string name() const { return "energy accumulator"; }

        /// \brief  return the energy
        double energy() const { return _energy; }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  returns the variance of the block averages for the specified block length
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        double bvar(const int nblocks) const {

          // check for sane number of blocks
          assert( nblocks > 0 );

          // get block length (last block will be longer if division does not result in an integer)
          const int bl = _history_e.size() / nblocks;

          // compute averages of each block on each process
          formic::Vector<double,int> avgs(nblocks, 0.0);
          for (int i = 0; i < nblocks; i++) {
            const int start = i * bl;
            const int end = ( i == nblocks-1 ? _history_e.size() : (i+1) * bl );
            avgs.at(i) = std::accumulate( _history_e.begin() + start, _history_e.begin() + end, 0.0) / double(end-start);
          }

          // compute average of each block across all processes
          formic::Vector<double,int> full_avgs(nblocks, 0.0);
          formic::mpi::reduce(&avgs.at(0), &full_avgs.at(0), nblocks, MPI::SUM);
          full_avgs /= double(formic::mpi::size());

          // compute overall average
          const double avg = std::accumulate(full_avgs.begin(), full_avgs.end(), 0.0) / double(nblocks);

          // compute the variance of the different blocks' averages
          double var = 0.0;
          for (int i = 0; i < nblocks; i++) {
            const double x = avg - full_avgs.at(i);
            var += x*x;
          }
          var /= double(nblocks);

          // return the variance on all processes
          formic::mpi::bcast(var);
          return var;

        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  performs a recursive blocking analysis of the statistical energy error
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        double recursive_blocking(formic::Output<std::ofstream> & fout, const bool print = true) {

          // make sure the sample was not an exact sample
          if ( _exact_sampling )
            throw formic::Exception("recursive blocking error analysis should not be performed for exact sampling");

          // get number of processors
          const double nproc = double(formic::mpi::size());

          // get copies of weights, numerators, and denominators
          std::vector<double> wtv(_history_p);
          std::vector<double> nmv(_history_ev);
          std::vector<double> dnv(_history_v);

          //// get vector to store sample counts and variances
          //std::vector<double> to_reduce;

          // get vector to estimates of the errors at different blocking levels
          std::vector<double> estimates;

          // each process gets the sample count and variance for each level of blocking
          int sample_end = nmv.size();
          for (int blk_lvl = 0; sample_end >= 200; blk_lvl++) {

            // compute and print error estimate for this level of blocking
            double eng, var;
            formic::mpi_unbiased_ratio_of_means(sample_end, &wtv.at(0), &nmv.at(0), &dnv.at(0), eng, var);
            estimates.push_back( std::sqrt( var / sample_end / nproc ) );
            if ( formic::mpi::rank() == 0 && print )
              fout << boost::format("  error estimate for blocks of size 2^(%2i) = %20.12f") % (blk_lvl+1) % (*estimates.rbegin()) << std::endl;

            // create new sample set by averaging adjacent elements of previous sample
            int j = 0;
            //to_reduce.push_back(0.0); // accumulation for variance of new sample
            for (int i = 0; i + 1 < sample_end; i += 2, j += 1) {
              wtv[j] = wtv[i] + wtv[i+1];
              nmv[j] = ( wtv[i] * nmv[i] + wtv[i+1] * nmv[i+1] ) / ( wtv[i] + wtv[i+1] );
              dnv[j] = ( wtv[i] * dnv[i] + wtv[i+1] * dnv[i+1] ) / ( wtv[i] + wtv[i+1] );
              //const double diff = _energy - nmv[j];
              //*to_reduce.rbegin() += diff * diff;
            }

            //// record number of elements from this process in new sample
            //to_reduce.push_back(double(j));

            // move the end-of-sample value to the end of the new, smaller sample set
            sample_end = j;

          }

          //// reduce the variance accumulations and sample counts to the root process
          //std::vector<double> reduced(to_reduce.size(), 0.0);
          //formic::mpi::reduce(&to_reduce.at(0), &reduced.at(0), to_reduce.size(), MPI::SUM);

//          // root process computes and prints error estimates
//          double err_est;
//          if ( formic::mpi::rank() == 0 ) {
//
//            // compute and print statistical error estimates for each level of blocking
//            int blk_lvl = 0;
//            for ( ; 2*blk_lvl+1 < reduced.size(); blk_lvl++) {
//              const double err = std::sqrt( reduced.at(2*blk_lvl) ) / reduced.at(2*blk_lvl+1);
//              fout << boost::format("  error estimate for blocks of size 2^(%2i) = %20.12f") % (blk_lvl+1) % err << std::endl;
//              reduced.at(blk_lvl) = err; // save the error estimate for this block level
//            }
//            fout << std::endl;
//
//            // compute overall error estimate as the average of that of the last few blocks
//            if ( blk_lvl < 1 )
//              throw formic::Exception("no blocks in energy's statistical error analysis");
//            const int n_to_avg = std::min(4, blk_lvl);
//            reduced.resize(blk_lvl);
//            err_est = std::accumulate(reduced.rbegin(), reduced.rbegin()+n_to_avg, 0.0) / double(n_to_avg);
//
//          }

          // compute overall error estimate as the average of that of the last few blocks
          if ( estimates.size() < 1 )
            throw formic::Exception("no blocks in energy's statistical error analysis");
          const int n_to_avg = std::min(4, int(estimates.size()));
          _serr = std::accumulate(estimates.rbegin(), estimates.rbegin()+n_to_avg, 0.0) / double(n_to_avg);

          // return the overall error estimate
          return _serr;

        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  Prints some statistics about the most recent sample.
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        void print_statistics(formic::Output<std::ofstream> & fout) {

          double red_target = 0.0;
          
          // compute the mean and variance
          double mean = _energy;
          //assert( _history_e.size() == _history_p.size() );
          //double mean = 0.0;
          //for (int i = 0; i < _history_p.size(); i++) {
          //  assert( _history_p.at(i) >= 0.0 );
          //  mean += _history_p.at(i) * _history_e.at(i);
          //}
          ////double mean = std::accumulate(_history.begin(), _history.end(), 0.0) / _history.size();
          //formic::mpi::allreduce(&mean, &red_target, 1, MPI::SUM);
          //mean = red_target;

          //// compute the variance
          //double var = _variance;
          //double var = 0.0;
          //for (int i = 0; i < _history_p.size(); i++) {
          //  assert( _history_p.at(i) >= 0.0 );
          //  var += _history_p.at(i) * formic::square_norm(mean - _history_e.at(i));
          //}
          ////typedef typename Accumulator<S>::VarBinOp VBOp;
          ////double var = std::accumulate(_history.begin(), _history.end(), 0.0, VBOp(mean)) / _history.size();
          //formic::mpi::allreduce(&var, &red_target, 1, MPI::SUM);
          //var = red_target;

          //fout << std::endl;
          //fout << boost::format("printing every sample:") << std::endl;
          //for (int i = 0; i < _history_e.size(); i++)
          //  fout << boost::format("  %20.12f  %20.12f  %20.12f") % _history_p.at(i) % _history_v.at(i) % _history_e.at(i) << std::endl;
          //fout << std::endl;

          //fout << boost::format("printing every 2nd sample:") << std::endl;
          //for (int i = 0; i < _history_e.size(); i+=2)
          //  fout << boost::format("  %20.12f  %20.12f  %20.12f") % _history_p.at(i) % _history_v.at(i) % _history_e.at(i) << std::endl;
          //fout << std::endl;

          //fout << boost::format("printing every 4th sample:") << std::endl;
          //for (int i = 0; i < _history_e.size(); i+=4)
          //  fout << boost::format("  %20.12f  %20.12f  %20.12f") % _history_p.at(i) % _history_v.at(i) % _history_e.at(i) << std::endl;
          //fout << std::endl;

          //fout << boost::format("printing every 8th sample:") << std::endl;
          //for (int i = 0; i < _history_e.size(); i+=8)
          //  fout << boost::format("  %20.12f  %20.12f  %20.12f") % _history_p.at(i) % _history_v.at(i) % _history_e.at(i) << std::endl;
          //fout << std::endl;

          //// compute the 10-block variance
          //const double bvar10 = ( _exact_sampling ? 0.0 : this->bvar(10) );
          //const double bvar100 = ( _exact_sampling ? 0.0 : this->bvar(100) );

          //// compute the effective number of samples for an independent sample (one measure of the autocorrelation time)
          //const double autocorr10 = ( bvar10 / _variance ) * formic::mpi::size() * _history_e.size() / 10.0;
          //const double autocorr100 = ( bvar100 / _variance ) * formic::mpi::size() * _history_e.size() / 100.0;

          //// compute the expected statistical error bars (one sigma) based on the two blocking analysis
          //const double err10 = std::sqrt(bvar10/10.0);
          //const double err100 = std::sqrt(bvar100/100.0);

          if ( formic::mpi::rank() == 0 ) {
            fout << boost::format("energy accumulation statistics:") << std::endl;
            fout << std::endl;
          }

          // print out recursive blocking information
          const double err_rec = ( _exact_sampling ? 0.0 : this->recursive_blocking(fout) );

          if ( formic::mpi::rank() == 0 ) {
            //fout << boost::format("energy accumulation statistics on process %i:") % formic::mpi::rank() << std::endl;
            //fout << std::endl;
            //fout << boost::format("   acceptance ratio = %20.12f") % (this->_naccept / this->_total_weight) << std::endl;
            fout << std::endl;
            fout << boost::format("               mean = %20.12f") % mean << std::endl;
            if ( _exact_sampling ) {
              fout << boost::format("           stat err = %20s") % "N/A" << std::endl;
              fout << boost::format("           autocorr = %20s") % "N/A" << std::endl;
              //fout << boost::format("            std dev = %20s") % "N/A" << std::endl;
              //fout << boost::format("           variance = %20s") % "N/A" << std::endl;
              //fout << boost::format(" recursive stat err = %20s") % "N/A" << std::endl;
              //fout << boost::format(" 10-block  stat err = %20s") % "N/A" << std::endl;
              //fout << boost::format(" 100-block stat err = %20s") % "N/A" << std::endl;
            } else {
              fout << boost::format("           stat err = %20.12f") % err_rec << std::endl;
              fout << boost::format("           autocorr = %20.12f") % ( err_rec * err_rec * formic::mpi::size() * _history_e.size() / _variance ) << std::endl;
              //fout << boost::format(" recursive stat err = %20.12f") % err_rec << std::endl;
              //fout << boost::format(" 10-block  stat err = %20.12f") % err10 << std::endl;
              //fout << boost::format(" 100-block stat err = %20.12f") % err100 << std::endl;
            }
            fout << boost::format("            std dev = %20.12f") % std::sqrt(_variance) << std::endl;
            fout << boost::format("           variance = %20.12f") % _variance << std::endl;
            //fout << boost::format("  10-block autocorr = %20.12f") % autocorr10 << std::endl;
            //fout << boost::format(" 100-block autocorr = %20.12f") % autocorr100 << std::endl;
            fout << std::endl;
          }

          //// compute the autocorrelation for short times
          //if ( !_exact_sampling ) {
          //  typedef typename Accumulator<S>::AutoCorrBinOp ACBOp;
          //  std::vector<double> ac(_max_ac_length, 0.0);
          //  for (int i = 1; i <= ac.size(); i++)
          //    ac.at(i-1) = std::accumulate(_history_e.begin()+i, _history_e.end(), 0.0, ACBOp(mean, i)) / ( _history_e.size() - i ) / _variance;
          //  if ( ac.size() > 0 ) {
          //    std::vector<double> ac_red(ac);
          //    formic::mpi::allreduce(&ac_red.at(0), &ac.at(0), ac.size(), MPI::SUM);
          //    formic::xscal(ac.size(), 1.0 / formic::mpi::size(), &ac.at(0), 1);
          //  }
          //  if ( formic::mpi::rank() == 0 && ac.size() > 0 ) {
          //    for (int i = 1; i <= ac.size(); i++)
          //      fout << boost::format("autocorrelation %3i = %20.12f") % i % ac.at(i-1) << std::endl;
          //    fout << std::endl;
          //  }
          //}

        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief Accumulate the local energy.
        ///
        /// \param[in]      userinp  the user's input options
        /// \param[in]      w        the weight to apply to this configuration's contribution
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        void child_accumulate(const formic::InputBase & userinp, const double w) {

          //std::cout << boost::format("vf --> %6i        gf -->  %24.12e") % _vf->get_coeff().exponent() % _gf->get_coeff().get() << std::endl;
          //formic::of << boost::format(" %20.12f <-- weight for energy acc") % w << std::endl;

          if ( w < 0.0 )
            throw formic::Exception("unexpected negative weight in local energy accumulator");

          if ( w == 0.0 )
            return;

          const double current_vgs = ( _vf == _gf  ?  1.0  :  formic::square_norm( ( _vf->get_coeff() / _gf->get_coeff() ).get() ) );
          const double current_eng = formic::real(_vf->local_energy());

          // print energy if requested
          if ( _print_local_e_during_sampling && formic::mpi::rank() == 0 )
            formic::of << boost::format(" %20.12f%s") % current_eng % ( _exact_sampling ? "" : "\n" );

          // exact sampling records just one probability and associated value/guiding ratio and local energy
          if ( _exact_sampling ) {

            // update the weight histories
            _history_p.push_back(w);

            // update the value/guiding ratio history
            _history_v.push_back(current_vgs);

            // update the energy history
            _history_e.push_back(current_eng);

          // stochastic sampling splits repeated samples of the same occupation number vector into multiple samples
          } else {

            double ww = 0.0;
            for ( ; w - ww > 1.0; ww += 1.0) {

              // update the weight histories
              _history_p.push_back(1.0);

              // update the value/guiding ratio history
              _history_v.push_back(current_vgs);

              // update the energy history
              _history_e.push_back(current_eng);

            }

            if ( w - ww <= 0.0 )
              throw formic::Exception("non-positive weight left over in local energy accumulator");

            // update the weight histories
            _history_p.push_back(w-ww);

            // update the value/guiding ratio history
            _history_v.push_back(current_vgs);

            // update the energy history
            _history_e.push_back(current_eng);

          }

        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief finalizes the accumulation (this may include for example communicating the results)
        ///
        /// \param[in]      userinp  the user's input options
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        void child_finalize(const formic::InputBase & userinp) {

          // update local energies using value/guiding ratios
          _history_ev.resize(_history_e.size());
          for (int i = 0; i < _history_e.size(); i++)
            _history_ev.at(i) = _history_e.at(i) * _history_v.at(i);

          // for a random sample, compute unbiased estimates of the energy and its variance
          if ( !_exact_sampling ) {
            formic::mpi_unbiased_ratio_of_means(_history_p.size(), &_history_p.at(0), &_history_ev.at(0), &_history_v.at(0), _energy, _variance);

          // for an exact sample, compute the energy and variance with the usual formulae
          } else {
            double x[3];
            x[0] = 0.0;
            x[1] = 0.0;
            x[2] = 0.0;
            for (int i = 0; i < _history_p.size(); i++) {
              x[0] += _history_p.at(i) * _history_v.at(i); // sum of probabilities
              x[1] += _history_p.at(i) * _history_ev.at(i); // weighted sum of local energies
              x[2] += _history_p.at(i) * _history_ev.at(i) * _history_e.at(i); // weighted sum of squares of local energy
            }
            double y[3];
            formic::mpi::allreduce(&x[0], &y[0], 3, MPI::SUM);
            y[1] /= y[0]; // average of local energy
            y[2] /= y[0]; // average of square of local energy
            _energy = y[1];
            _variance = y[2] - y[1] * y[1];
          }

          // we have not yet calculated statistical error, so set it to negative to indicate this
          _serr = -1.0;

        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  resets the accumulator
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        void child_reset(const formic::Configuration & config) {
          _energy = 0.0;
          _variance = 0.0;
          _serr = -1.0;
          _history_p.clear();
          _history_v.clear();
          _history_e.clear();
          _history_ev.clear();
        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  return the variance
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        double variance() const {
          return _variance;
        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief  return the statistical error
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        double serr() {
          if ( _exact_sampling )
            return 0.0;
          if ( _serr < 0.0 )
            this->recursive_blocking(formic::of, false);
          return _serr;
        }

    };

  }

}

#endif
